import streamlit as st
import requests
from PIL import Image
from io import BytesIO

# FastAPI 服务的 URL
API_URL = "http://192.168.1.90:8000/generate-text"
CLEAR_HISTORY_URL = "http://192.168.1.90:8000/clear_history"

# 初始化 session state
if "conversation_history" not in st.session_state:
    st.session_state.conversation_history = []

# Streamlit 应用标题
st.title("多模态问答系统")

# 上传图片
uploaded_image = st.file_uploader("上传图片", type=["jpg", "jpeg", "png"])

# 输入文本
user_input = st.text_input("输入你的问题")

# 提交按钮
if st.button("提交"):
    if uploaded_image is not None and user_input:
        # 将图片转换为 PIL.Image 对象
        pil_image = Image.open(uploaded_image)

        # 将图片转换为字节流
        image_bytes = BytesIO()
        pil_image.save(image_bytes, format="PNG")
        image_bytes.seek(0)

        # 构造请求数据
        files = {"image": ("image.png", image_bytes, "image/png")}
        data = {"text": user_input}

        # 发送请求到 FastAPI 服务
        response = requests.post(API_URL, files=files, data=data)

        if response.status_code == 200:
            # 获取模型生成的文本
            output_text = response.json()["output_text"]

            # 更新对话历史
            st.session_state.conversation_history.append({"role": "user", "content": user_input})
            st.session_state.conversation_history.append({"role": "assistant", "content": output_text})

            # 显示对话历史
            # st.write("对话历史:")
            for msg in st.session_state.conversation_history:
                if msg["role"] == "user":
                    st.write(f"用户: {msg['content']}")
                else:
                    st.write(f"助手: {msg['content'][0]}")
        else:
            st.error("请求失败，请稍后再试。")
    else:
        st.warning("请上传图片并输入问题。")

# 清除记录按钮
if st.button("清除记录"):
    # 调用 FastAPI 的 /clear_history 接口
    response = requests.post(CLEAR_HISTORY_URL)
    if response.status_code == 200:
        st.session_state.conversation_history = []
        st.write("对话记录已清除。")
    else:
        st.error("清除记录失败，请稍后再试。")
